import time
from base64 import b64decode
from binascii import a2b_hex, b2a_hex
from concurrent.futures import Future, ProcessPoolExecutor
from json import loads
from os import listdir, path
from struct import unpack

from Crypto.Cipher import AES
from mutagen.flac import FLAC, Picture
from mutagen.id3 import APIC, ID3
from PySide6.QtCore import QObject, Signal
from PySide6.QtGui import QColor


##############################################################################################################
# NcmDecrypt
##############################################################################################################
class NotNcmError(Exception):
    pass


class TransformResult:
    def __init__(self, msg, color) -> None:
        self.msg = msg
        self.color = color


unpad = lambda s: s[0 : -(s[-1] if type(s[-1]) == int else ord(s[-1]))]


def load_ncm(file_path):
    with open(file_path, "rb") as ncm_file:
        header = ncm_file.read(8)
        if b2a_hex(header) != b"4354454e4644414d":
            raise NotNcmError(f"{file_path} is not ncm file.")
        ncm_file.seek(2, 1)
        key_length = unpack("<I", ncm_file.read(4))[0]
        key_data = ncm_file.read(key_length)
        meta_length = unpack("<I", ncm_file.read(4))[0]
        meta_data = ncm_file.read(meta_length)
        crc32 = unpack("<I", ncm_file.read(4))[0]
        ncm_file.seek(5, 1)
        image_size = unpack("<I", ncm_file.read(4))[0]
        image_data = ncm_file.read(image_size)
        audio_data = bytearray(ncm_file.read())
    return {
        "key_data": key_decode(key_data),
        "meta_data": meta_decode(meta_data),
        "image_data": image_data,
        "audio_data": audio_data,
        "crc32": crc32,
    }


def key_decode(key_data):
    core_key = a2b_hex("687A4852416D736F356B496E62617857")
    key_data = bytearray(key_data)
    for i in range(0, len(key_data)):
        key_data[i] ^= 0x64
    cryptor = AES.new(core_key, AES.MODE_ECB)
    key_data = unpad(cryptor.decrypt(key_data))[17:]
    return key_data


def meta_decode(meta_data):
    meta_key = a2b_hex("2331346C6A6B5F215C5D2630553C2728")
    meta_data = bytearray(meta_data)
    for i in range(0, len(meta_data)):
        meta_data[i] ^= 0x63
    meta_data = b64decode(meta_data[22:])
    cryptor = AES.new(meta_key, AES.MODE_ECB)
    meta_data = unpad(cryptor.decrypt(meta_data)).decode("utf-8")[6:]
    return loads(meta_data)


def get_keybox(key_data):
    box = [i for i in range(256)]
    key_data_len = len(key_data)
    j = 0
    for i in range(256):
        j = (box[i] + j + key_data[i % key_data_len]) & 0xFF
        box[i], box[j] = box[j], box[i]
    boxmap = []
    for i in range(256):
        i = (i + 1) & 0xFF
        si = box[i]
        sj = box[(i + si) & 0xFF]
        boxmap.append(box[(si + sj) & 0xFF])
    return boxmap


def transform(ncm_file, out_path):
    try:
        ncm = load_ncm(ncm_file)
        audio_data = ncm.get("audio_data")
        key_box = get_keybox(ncm.get("key_data"))
        for i in range(len(audio_data)):
            audio_data[i] ^= key_box[i & 0xFF]
        format: str = ncm.get("meta_data").get("format")
        image_data = ncm.get("image_data")
        file_name, extension = path.splitext(path.basename(ncm_file))
        out_file = path.join(out_path, f"{file_name}.{format}")
        with open(out_file, "wb") as f:
            f.write(bytes(audio_data))

        if format.lower() == "flac":
            audio = FLAC(out_file)
            artwork = Picture()
            artwork.type = 3
            artwork.mime = "image/jpeg"
            artwork.desc = "Cover"
            artwork.data = image_data
            audio.clear_pictures()
            audio.add_picture(artwork)
            audio.save()

        if format.lower() == "mp3":
            audio = ID3(out_file)
            audio.update_to_v23()
            apic = APIC(
                encoding=0, mime="image/jpeg", type=3, desc="Cover", data=image_data
            )
            audio.add(apic)
            audio.save()
    except NotNcmError:
        return TransformResult(f"{path.basename(ncm_file)} 无法识别", QColor(198, 15, 15))
    except Exception as error:
        return TransformResult(f"转换错误: {error}", QColor(198, 15, 15))
    return TransformResult(f"{path.basename(ncm_file)} 转换完成", QColor(0, 0, 0))


##############################################################################################################


class TransformObject(QObject):
    startSig = Signal(str)
    resultSig = Signal(TransformResult)
    finishSig = Signal(str, QColor)

    def set_attr(self, ncm_path: str, out_path: str, max_workers: int):
        self.ncm_path = ncm_path
        self.out_path = out_path
        self.max_workers = max_workers

    def transform(self):
        ncm_files = self._get_ncm_converted(self.ncm_path, self.out_path)
        if not ncm_files:
            self.finishSig.emit("没有需要转换的ncm文件", QColor(198, 15, 15))
            return

        max_workers = min(self.max_workers, len(ncm_files))

        self.startSig.emit(f"开始转换{len(ncm_files)}个文件")

        start_time = time.time()

        process_pool = ProcessPoolExecutor(max_workers)

        task_list: list[Future] = []
        for ncm_file in ncm_files:
            task = process_pool.submit(transform, ncm_file, self.out_path)
            task_list.append(task)

        for task in task_list:
            if task.result():
                self.resultSig.emit(task.result())

        end_time = time.time()
        self.finishSig.emit(
            f"全部转换完成，耗时{int(end_time-start_time)}秒",
            QColor(15, 198, 15),
        )

    def _get_ncm_converted(self, ncm_path: str, out_path: str):
        exists_out_files = [
            path.splitext(file)[0].lower()
            for file in listdir(out_path)
            if path.isfile(path.join(out_path, file))
            and path.splitext(file)[1] != ".ncm"
        ]
        return [
            path.join(ncm_path, name)
            for name in listdir(ncm_path)
            if path.splitext(name)[1] == ".ncm"
            and path.splitext(name)[0].lower() not in exists_out_files
        ]
